"""
This module lets you run a Python script and track which modules it imported.
"""

import sys
import functools
import inspect
import modulefinder
import __builtin__

# modefinder works well, but it is unfortunately too slow -- often an
# order of magnitude slower than actually running the script.  Instead,
# we use the old-style import hook mechanism.
use_modulefinder = False

def run_script_and_get_imported_modules(script_path, *args, **kwargs):
    """Call execfile on the given script, and return a set containing the
    absolute module names of the modules that were imported by the script.

    The result includes modules that have already been imported and modules
    that were imported indirectly by importing another module.

    Any additional parameters will be passed to execfile.  This function does
    not trap any exceptions raised by execfile.
    """
    function = (_run_script_and_get_imported_modules_with_modulefinder
        if use_modulefinder else
        _run_script_and_get_imported_modules_with_hook)
    module_names = function(script_path, *args, **kwargs)

    # We don't want to include __main__ in the result.
    if "__main__" in module_names:
        module_names.remove("__main__")
    return module_names

def _run_script_and_get_imported_modules_with_modulefinder(
        script_path, *args, **kwargs):
    # modulefinder uses static analysis to recursively determine the imported
    # modules.
    module_finder = modulefinder.ModuleFinder()
    module_finder.run_script(script_path)

    execfile(script_path, *args, **kwargs)

    return set(module_finder.modules.keys())

def _run_script_and_get_imported_modules_with_hook(
        script_path, *args, **kwargs):
    with _ImportTracker() as import_tracker:
        execfile(script_path, *args, **kwargs)

    return import_tracker.imported_module_names

class _ImportTracker(object):
    """A context guard that replaces the builtin __import__ with a custom
    function that tracks which modules are imported.
    """
    def __init__(self):
        self.imported_module_names = set()

    def __enter__(self):
        # Note that we can't use PEP 302 style import hooks because they are
        # not invoked when the module is already in sys.modules, and we want
        # to detect scripts that import modules that may have already been
        # imported.
        self.saved_builtin_import = __builtin__.__import__

        __builtin__.__import__ = functools.wraps(__builtin__.__import__)(
            lambda name, globals=None, locals=None, fromlist=None, level=-1:
                self._replacement_import(
                    name, globals, locals, fromlist, level))

        return self

    def __exit__(self, type, value, traceback):
        __builtin__.__import__ = self.saved_builtin_import

    def _replacement_import(
            self, name, globals_dict, locals_dict, from_seq, level):
        """A replacement for the builtin __import__ that delegates to the
        standard builtin.

        Note that __import__ doesn't always return the module that was
        imported, depending on whether packages are involved.
        """
        # If globals was not provided and the user attempts a relative import
        # from within a package, the standard behaviour of __import__ is to
        # not find the module.  So, we don't need to bother to try to find
        # the caller frame's globals dict if globals_dict is None.

        return_module = self.saved_builtin_import(
            name, globals_dict, locals_dict, from_seq, level)

        module_name = self._determine_imported_module(name, globals_dict, level)
        self.imported_module_names.add(module_name)

        # It's possible to write "from mod1 import mod2, func1", importing
        # both a module (mod2) from a module (mod1) as well as a non-module
        # (func1).  In this case we'll have already added mod1 as an imported
        # module, and we need to check if any imported names are also modules.
        if from_seq is not None:
            for attrib_name in from_seq:
                submodule_name = "%s.%s" % (module_name, attrib_name)
                if submodule_name in sys.modules:
                    self.imported_module_names.add(submodule_name)

        return return_module

    def _determine_imported_module(
            self, requested_name, globals_dict, level):
        # Note that when from_seq is empty and we're importing a dotted module
        # name with an absolute import, __import__ returns the toplevel package
        # module instead of the module being imported.  However, when doing a
        # relative import from within a package, __import__ returns the module
        # being imported.

        # If level is 0 or the import is not done from within a package then
        # the name passed to __import__ must be an absolute module name.
        caller_package = self._get_caller_package(globals_dict)
        if level == 0 or caller_package is None:
            assert(requested_name in sys.modules)
            return requested_name

        # If a positive level is given then go up from the caller's package
        # and find the module there.  For example, if a.b.c runs
        # "from .. import d" then level=1, caller_package is a.b, and we want
        # to import a.b.d.
        if level > 0:
            package = self._get_parent_package(caller_package, level - 1)
            module_name = (package if requested_name == ""
                else "%s.%s" % (package, requested_name))
            assert(module_name in sys.modules)
            return module_name

        # Check if the requested name works for a relative import.  Otherwise,
        # it must be an absolute import.
        assert(level < 0)
        module_name = "%s.%s" % (caller_package, requested_name)
        return (module_name if module_name in sys.modules else requested_name)

    def _get_caller_package(self, globals_dict):
        if globals_dict is None:
            return None

        # Although Python's import.c uses __path__ to determine the package,
        # we purposely ignore it since we want absolute module names, not file
        # paths.

        # If __package__ was set and isn't None or an empty string then use it.
        package = globals_dict.get("__package__")
        if package:
            return package

        # If the caller has a __path__ variable then the caller is a
        # __init__.py file in a package, and __name__ is the name of the
        # package.
        if "__path__" in globals_dict:
            return globals_dict["__name__"]

        # This is a normal module.
        return self._get_parent_package(globals_dict["__name__"])

    def _get_parent_package(self, package, level=1):
        for i in range(level):
            package = package.rpartition(".")[0]
        return (None if package == "" else package)

